import numpy as np
import time
from .top_n import top_n


class Genetic:
    def __init__(self, num_bodies=20, num_children=60, num_generations=200, mutation_rate=0.1):
        """
        创建一个遗传算法对象
        :param num_bodies: 被遗传的生物体的数量
        :param num_children: 每代生成的子代数量
        :param num_generations: 遗传多少代
        :param mutation_rate: 变异率，发生变异的概率
        """
        self._num_bodies = num_bodies
        self._num_children = num_children
        self._num_generations = num_generations
        self._mutation_rate = mutation_rate

    def get_random_dnas(self, num):
        raise Exception('get_random_dna() is not defined')

    def dnas2bodies(self, dna_s):
        raise Exception('dna2body() is not defined')

    def get_scores(self, bodies):
        raise Exception('get_score() is not defined')

    def mutate(self, dna_s, new_dna=None):  # 变异
        """
        变异
        :param dna_s:
        :param new_dna:
        :return:
        """
        if type(dna_s) in (list, tuple, set, dict):
            new_dna = self.get_random_dnas(1)  # [dna_size]
            return [self.mutate(dna, new_dna[i][0]) for i, dna in enumerate(dna_s)]
        if new_dna is None:
            new_dna = next(iter(self.get_random_dnas(1)))  # [dna_size]
        # dna_s: [n, dna_size]
        r = np.random.uniform(0, 1, dna_s.shape) <= self.get_mutation_rate()  # [n, dna_size]
        return (1 - r) * dna_s + r * new_dna

    def hybridize(self, dna_s, index_pairs):  # 杂交
        """
        杂交
        :param dna_s: 基因序列或者基因序列的列表
        :param left: 杂交的左边dna索引列表，len(left)==len(right)
        :param right: 杂交的右边dna索引列表，len(left)==len(right)
        :return: 杂交基因
        """
        if type(dna_s) in (list, tuple, set, dict):
            return [self.hybridize(dna, index_pairs) for dna in dna_s]
        dna_s = dna_s[index_pairs]  # [n, 2, dna_size]
        shape = (len(index_pairs), 1, *dna_s.shape[2:])  # [n, 1, dna_size]
        r = np.random.randint(0, 2, shape)  # [n, 1, dna_size]
        r = np.concatenate((r, 1-r), 1)  # [n, 2, dna_size]
        return np.sum(dna_s * r, 1)   # [n, dna_size]

    def get_num_bodies(self):
        return self._num_bodies

    def get_num_generations(self):
        return self._num_generations

    def get_num_children(self):
        return self._num_children

    def get_mutation_rate(self):
        return self._mutation_rate

    def get_best_body(self):
        dna_s = self.get_random_dnas(self.get_num_children())
        is_iterable = type(dna_s) in (list, tuple, set, dict)
        for gen in range(self.get_num_generations()):
            dna_s, i = self.screen(dna_s)
            best = [[d[i]] for d in dna_s] if is_iterable else [dna_s[i]]
            self.before_generate(gen)
            dna_s = self.generate(dna_s)
        best = self.dnas2bodies(best)
        return [d[0] for d in best] if is_iterable else best[0]

    def screen(self, dna_s):  # 基因筛选
        scores = self.get_scores(self.dnas2bodies(dna_s))

        def _compare(i, j):
            return scores[i] - scores[j]
        indices = top_n(range(len(scores)), self.get_num_bodies(), _compare)
        is_iterable = type(dna_s) in (list, tuple, set, dict)
        dna_s = [d[indices] for d in dna_s] if is_iterable \
                else dna_s[indices]
        scores = np.array(scores)[indices]
        self.after_screen(dna_s, scores)
        i = np.argmax(scores)
        return dna_s, i

    def after_screen(self, dnas, scores):
        print("max, min, avg score: %.5f,\t%.5f,\t%.5f" % (np.max(scores),
              np.min(scores), np.mean(scores)))

    def generate(self, dna_s):  # 生成下一代
        num = self.get_num_children()
        is_iterable = type(dna_s) in (tuple, list, set, dict)
        size = len(dna_s[0] if is_iterable else dna_s)
        left = np.random.randint(0, size-1, [num])
        right = np.random.randint(left+1, size, [num])
        index_pairs = np.transpose((left, right))  # [n, 2]
        dna_s2 = self.mutate(self.hybridize(dna_s, index_pairs))
        return [np.concatenate((d, d2), 0) for d, d2 in zip(dna_s, dna_s2)] \
                   if is_iterable else \
               np.concatenate((dna_s, dna_s2), axis=0)

    def before_generate(self, gen):
        t = time.time()
        t = time.strftime('%m-%d %H:%M:%S', time.localtime(t))
        print('%s: generation: %d/%d' % (t, gen + 1, self.get_num_generations()))
